Let's reproduce GPT-2(124M)

2025年6月28日

11:07

• pytorch张量默认都是fp32的，即32位浮点数，对于深度学习训练来说，不需要这么高的精度。

 

• int8用于推理，不用于训练，因为int8基本具有均匀的间距分布，我们需要浮点数，来满足训练期间，激活值和权重的正态分布。

 

• GPU内存的带宽问题，tensor cores大部分时间并没有在计算，而是处于空闲状态，在等待数据。所以降低tensor的精度，不仅可以减少显存使用，还能使数据更快地传输。
• 什么是tensor core?tensor core是GPU架构中的一个指令，它的作用基本上就是一个4*4矩阵乘法，任何需要矩阵乘法的操作都被 分解成这种小的4*4乘法。
• 如果要统计在GPU上的运行时间，需要用torch.cuda.synchronize()，因为当CPU运行时，它只是在GPU上调度工作，给GPU安排 一些工作，安排后CPU继续运行，所以需要等待GPU完成工作后，才能计算时间。
• 在训练时，每秒处理的token数量是我们真正关心的客观指标。使用torch.set_float32_matual_precision('high')设置，tensor cores在计算时使用tf32进行运算而不是fp32，理论上每秒处理的token数即吞吐量会快8倍，因为使用tf32计算的速度比fp32快8倍，但是实际上只快了3倍，这是因为只是在计算时使用了tf32，在tensor cores之外，这些数仍然是fp32，这些fp32数字通过内存系统进行传输，传输过程占用了很多时间。因此，尽管我们已经使乘法本身更快了，但我们仍然受限于内存带宽。
• 关于浮点数。即科学计数法。包含两部分，指数部分+小数部分。指数exponent决定了你能表示的数值范围，即range；小数部分manissa决定了数值的精确程度，即precision。

  BF16的range和FP32一样，不需要对梯度进行缩放，而FP16的range小于FP32，所以训练时需要梯度scaler。

• torch.compile加速，speedup mainly comes from reducing python overhead and GPU read/writes. 首先将神经网络编译成一个不涉及python解释器的单一对象，它明确知道要运行什么并直接运行；减少GPU read/writes，一个例子是GELU函数，如果不用torch.compile，计算GELU激活值时，要先将输入x从显存移到GPU缓存中，计算torch.pow(x,3)，然后再将计算结果移回显存中，再将结果从显存移到GPU缓存中，计算torch.tanh(x)….（GELU的等效分解公式），，，，，这样大量时间浪费在了传输上，但torch.compile能看到整个的计算代码，它能意识到这些操作都是element-wise操作，它会一次性将x移到GPU缓存中，并进行所有的操作，再把结果移到显存中，不会进行来回地传输。但是有些操作是torch.compile无法发现识别的，需要自定义，如flash attention.
• cuda中的kernels使用block tiles，即一小块一小块，这些block tiles是2的N次方，如64,32，也就是说，计算是以64或32这样的块进行的。所以你的输入也最好是32或64这些2的指数，否则在计算时，会启动一些额外的块来计算剩余部分，导致计算低效。例如GPT2(config={vocab_size:50057})，这个数不是2的指数，将其增加一些，变为50304，虽然计算量增加了，但是训练时间反而缩短了，就是这个原因。所以填充你的输入，使其变为2的倍数，有时会加速计算。
• torch.nn.utils.clip_grad_norm(model.parameters(),1.0)，我一直以为是对每个参数的梯度绝对值进行裁剪，其实是参数向量的范数，也就是对参数向量的L2范数（可以看作向量的长度） sqrt(w1^2 +w2^2 +… +wn^2)进行裁剪，具体怎么裁剪？所有参数等比例缩小？使用梯度裁剪的原因：在训练时可能遇到某个噪音比较大的或者数据质量很差的batch，在这个batch上的loss很高，loss高的话梯度也会大，会shock你的优化过程，通过梯度裁剪防止模型受到过大的冲击，对梯度进行上限控制，说白了，不希望参数一次性变动过大，而是慢慢地变化，是一种正则技巧。
• 线性预热余弦衰减学习率，学习率从接近0开始，在一段时间内线性提升，然后以余弦形式下降至某个你设定的学习率。在GPT3中是初始学习率的10%
• 
• 在训练的早期，即前几个step，模型尚处于不稳定的状态，模型主要学习的是忽略在训练集上不经常出现的token，此时每个样本的梯度是高度相关的，类似人的学习总结规律一样，刚开始只看了几个例子，此时人总结出的规律是bias的，不全面的。在训练后期，模型学会了所有简单的内容，真正的学习才开始，这时每个样本的梯度才变得不相关。
• 权重衰减weight decay。参数可以分为进行权重衰减的和不进行权重衰减的，不进行衰减的有bias，和其他一维的张量如layernorm的scale和bias参数，对这些一维参数进行衰减没有太大意义，主要是对参数矩阵和embedding矩阵进行衰减。权重衰减是一种正则，当每个参数的权重变得小时，会迫使模型使用更多的参数，而不是依赖某些参数。
• 关于梯度累积，其实是完全可以跟模拟的batch_size实现完全一致的梯度，只需要在计算loss时，进行缩放，乘以1/accumulation_steps，因为在小batch上的梯度累积等价于对小batch上loss的求和，而真正的大batch上的loss其实是loss/样本数量，差了个缩放因子。例子如下：

# super simple little MLP
net = torch.nn.Sequential(
    torch.nn.Linear(16, 32),
    torch.nn.GELU(),
    torch.nn.Linear(32, 1)
)
torch.random.manual_seed(42)

#这里是大batch, batch_size=4，一次性输入4个样本，进行loss.backward()

x = torch.randn(4, 16)
y = torch.randn(4, 1)
net.zero_grad()
yhat = net(x)
loss = torch.nn.functional.mse_loss(yhat, y)
loss.backward()
print(net[0].weight.grad.view(-1)[:10])

# the loss objective here is (due to readuction='mean')
# L = 1/4 * [
#            (y[0] - yhat[0])**2 +
#            (y[1] - yhat[1])**2 +
#            (y[2] - yhat[2])**2 +
#            (y[3] - yhat[3])**2
#           ]
# NOTE: 1/4!

 

# now let's do it with grad_accum_steps of 4, and B=1
# the loss objective here is different because
# accumulation in gradient <---> SUM in loss
# i.e. we instead get:
# L0 = 1/4(y[0] - yhat[0])**2
# L1 = 1/4(y[1] - yhat[1])**2
# L2 = 1/4(y[2] - yhat[2])**2
# L3 = 1/4(y[3] - yhat[3])**2
# L = L0 + L1 + L2 + L3
# NOTE: the "normalizer" of 1/4 is lost
net.zero_grad()

#这里是小batch, batch_size=1，每次输入1个样本，进行loss.backward，此时loss要进行缩放，即loss=loss/4

for i in range(4):
    yhat = net(x[i])
    loss = torch.nn.functional.mse_loss(yhat, y[i])
    loss = loss / 4 # <-- have to add back the "normalizer"!
    loss.backward()
print(net[0].weight.grad.view(-1)[:10])

 

 

import os

import math

import time

import inspect

from dataclasses import dataclass

import torch

import torch.nn as nn

from torch.nn import functional as F

from hellaswag import render_example, iterate_examples

# -----------------------------------------------------------------------------

 

class CausalSelfAttention(nn.Module):

#多头注意力并不复杂，就是多个头并行工作，它们的输出被简单地concat起来，形成了多头注意力的输出

#这里的实现跟lets build gpt中不同，不是先定义一个head，再定义Multihead，这里是直接实现多个head

    def __init__(self, config):

        super().__init__()

        assert config.n_embd % config.n_head == 0

        # key, query, value projections for all heads, but in a batch

        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) #这里用self.c_attn直接相当于多头的Wq,Wk,Wv权重矩阵。原来单个头的 self.query = nn.Linear(n_embd, head_size, bias=False)， self.key= nn.Linear(n_embd, head_size, bias=False), self.value = nn.Linear(n_embd, head_size, bias=False)，即单头的q,k,v的size都是B,T,head_size，多头的q,k,v的size都是B,T,n_embd，多头的q,k,v一共的size就是3倍的B,T,n_embd（因为q,k,v每个都是B,T,n_embd），也就是这里的self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)

        # output projection

        self.c_proj = nn.Linear(config.n_embd, config.n_embd)

        self.c_proj.NANOGPT_SCALE_INIT = 1

        # regularization

        self.n_head = config.n_head

        self.n_embd = config.n_embd

 

    def forward(self, x):

        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim

        # nh is "number of heads", hs is "head size", and C (number of channels) = nh * hs

        # e.g. in GPT-2 (124M), n_head=12, hs=64, so nh*hs=C=768 channels in the Transformer

        qkv = self.c_attn(x) #qkv先是一个整体，

        q, k, v = qkv.split(self.n_embd, dim=2) #在这里将q,k,v进行split，从而获得q,k,v，注意这里的q,k,v是多头的q,k,v，而不是单头的q,k,v，即它们的size是B,T,n_embd，而不是B,T,head_size。之前在let's build GPT的实现中k = self.key(x) ，q=self.query(x), v=self.value(x)，这种是分别得到单个头的q,k,v，然后再把n个head的q,k,v concat起来。

        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, headsize) 我们把头的数量nh变成了一个类似batch的维度，这样pytorch可以像处理batch一样并行地处理nh。

        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, headsize)

        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, headsize)

 

        y = F.scaled_dot_product_attention(q, k, v, is_causal=True) # flash attention，当执行这行代码时，pytorch会自动调用flash attention，

#flash attention相当于代替了下面这4行代码

        #att = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)

        #att = att.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)

        #att = F.softmax(att, dim=-1) # (B, T, T)

        #y    =  att@v

FlashAttention利用底层硬件的内存层次知识，例如GPU的内存层次结构，来提高计算速度和减少内存访问开销。 FlashAttention的核心原理是通过将输入分块并在每个块上执行注意力操作，从而减少对高带宽内存（HBM）的读写操作。具体而言，FlashAttention使用平铺和重计算等经典技术，将输入块从HBM加载到SRAM（快速缓存），在SRAM上执行注意力操作，并将结果更新回HBM。FlashAttention减少了内存读写量，从而实现了2-4倍的时钟时间加速。

 

FlashAttention将attention的四个操作，整合成单一的融合核。如果只看浮点运算数FLOPS，flash attention比原始的attention执行了更多的FLOPS。Flashattention使得attention得分矩阵即att=q@k.transpose(-2,-1)永远不会被具象化，也就不会被读写到HBM中，这是一个非常大的矩阵，B，T，T。

也就是说，原始的attention计算需要先算出q@k，然后再将结果转移到HBM显存，再将结果从HBM转移到SRAM，再计算mask，再进行转移，再计算softmax，再进行转移 ，再计算dropout,再进行转移，再计算attn@v。而flash attn一次性地转移，并计算最后的结果。

Flash attn使用了在线softmax技巧，（和移动平均类似？？），不需要得到softmax的所有元素，

Flash attn只是一个attention实现的重写，是一个更快的内核 ，并没有改变任何的计算，也不会提升算法效果。

 

 

        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection

        y = self.c_proj(y)

        return y

 

class MLP(nn.Module):

 

    def __init__(self, config):

        super().__init__()

        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd)

        self.gelu    = nn.GELU(approximate='tanh')

        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd) #MLP是两个线性层c_fc、c_proj中间夹了一个GELU激活函数

        self.c_proj.NANOGPT_SCALE_INIT = 1

 

    def forward(self, x):

        x = self.c_fc(x)

        x = self.gelu(x)

        x = self.c_proj(x)

        return x

 

class Block(nn.Module):

 

    def __init__(self, config):

        super().__init__()

        self.ln_1 = nn.LayerNorm(config.n_embd)

        self.attn = CausalSelfAttention(config)

        self.ln_2 = nn.LayerNorm(config.n_embd)

        self.mlp = MLP(config)

 

    def forward(self, x):

        x = x + self.attn(self.ln_1(x)) #预归一化，先经过layernorm，再经过attention。注意力是一个聚合函数，池化函数，加权和函数，是一个归约操作reduce operation，

        x = x + self.mlp(self.ln_2(x)) #先经过layernorm，再经过mlp。MLP发生在每个独立的token上，token之间没有交换信息，是一个map操作。所以Transformer可以看成一个不断进行map reduce的过程。

#残差连接，即加法，即在反向传播时将梯度平均地分配给它的两个分支

        return x

 

@dataclass

class GPTConfig:

    block_size: int = 1024 # max sequence length

    vocab_size: int = 50257 # number of tokens: 50,000 BPE merges + 256 bytes tokens + 1 <|endoftext|> token

    n_layer: int = 12 # number of layers

    n_head: int = 12 # number of heads

    n_embd: int = 768 # embedding dimension

 

官方GPT2的state_dict()，我们仿照这个state_dict进行模型构建

model_hf = GPT2LMHeadModel.from_pretrained("gpt2") # 124M

sd_hf = model_hf.state_dict()

 

for k, v in sd_hf.items():

    print(k, v.shape)

 

transformer.wte.weight torch.Size([50257, 768])
transformer.wpe.weight torch.Size([1024, 768])
transformer.h.0.ln_1.weight torch.Size([768])
transformer.h.0.ln_1.bias torch.Size([768])
transformer.h.0.attn.c_attn.weight torch.Size([768, 2304])
transformer.h.0.attn.c_attn.bias torch.Size([2304])
transformer.h.0.attn.c_proj.weight torch.Size([768, 768])
transformer.h.0.attn.c_proj.bias torch.Size([768])
transformer.h.0.ln_2.weight torch.Size([768])
transformer.h.0.ln_2.bias torch.Size([768])
transformer.h.0.mlp.c_fc.weight torch.Size([768, 3072])
transformer.h.0.mlp.c_fc.bias torch.Size([3072])
transformer.h.0.mlp.c_proj.weight torch.Size([3072, 768])
transformer.h.0.mlp.c_proj.bias torch.Size([768])
transformer.h.1.ln_1.weight torch.Size([768])
transformer.h.1.ln_1.bias torch.Size([768])
transformer.h.1.attn.c_attn.weight torch.Size([768, 2304])
transformer.h.1.attn.c_attn.bias torch.Size([2304])
transformer.h.1.attn.c_proj.weight torch.Size([768, 768])
transformer.h.1.attn.c_proj.bias torch.Size([768])
transformer.h.1.ln_2.weight torch.Size([768])
transformer.h.1.ln_2.bias torch.Size([768])
transformer.h.1.mlp.c_fc.weight torch.Size([768, 3072])
transformer.h.1.mlp.c_fc.bias torch.Size([3072])
transformer.h.1.mlp.c_proj.weight torch.Size([3072, 768])
transformer.h.1.mlp.c_proj.bias torch.Size([768])
transformer.h.2.ln_1.weight torch.Size([768])
transformer.h.2.ln_1.bias torch.Size([768])
transformer.h.2.attn.c_attn.weight torch.Size([768, 2304])
transformer.h.2.attn.c_attn.bias torch.Size([2304])
transformer.h.2.attn.c_proj.weight torch.Size([768, 768])
transformer.h.2.attn.c_proj.bias torch.Size([768])
transformer.h.2.ln_2.weight torch.Size([768])
transformer.h.2.ln_2.bias torch.Size([768])
transformer.h.2.mlp.c_fc.weight torch.Size([768, 3072])
transformer.h.2.mlp.c_fc.bias torch.Size([3072])
transformer.h.2.mlp.c_proj.weight torch.Size([3072, 768])
transformer.h.2.mlp.c_proj.bias torch.Size([768])
…….

transformer.ln_f.weight torch.Size([768])
transformer.ln_f.bias torch.Size([768])
lm_head.weight torch.Size([50257, 768])

注意，lm_head.weight和transformer.wte.weight的size是一样的，都是20257,768，实际上不只是大小一样，这两个其实是同一个tensor，即权重共享。原因：如果两个token在语义上是相似的，那么它们在嵌入空间中也应该是相邻的，即经过transformer.wte.weight后的embedding相邻，同样的，这两个相似语义的token经过lm_head.weight后的输出概率也应该是几乎相同的。即两个token语义相似<==>两个embedding相邻<==>两个输出logits相同。这两个矩阵具有相同的特性。因此进行权重共享。

 

 

class GPT(nn.Module):

 

    def __init__(self, config):

        super().__init__()

        self.config = config

 

        self.transformer = nn.ModuleDict(dict( #如上所示，官方GPT2的主模块的名称是transformer，是一个ModuleDict，可以像字典一样通过key来索引到子模块

            wte = nn.Embedding(config.vocab_size, config.n_embd), #token embedding

            wpe = nn.Embedding(config.block_size, config.n_embd), #位置embedding

            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),# transformer.h，是一个ModuleList，可以像列表一样用整数进行索引，h[0],h[1]…h[11]，每个元素是一个Blcok，Block的定义在上面

            ln_f = nn.LayerNorm(config.n_embd), #GPT2最后有一个layernorm层

        ))

        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

 

        # weight sharing scheme

        self.transformer.wte.weight = self.lm_head.weight

 

        # init params

        self.apply(self._init_weights)

 

    def _init_weights(self, module):

        if isinstance(module, nn.Linear):

            std = 0.02 #大概是1/sqrt(d_model)的大小，因为矩阵A*矩阵B后，矩阵A的方差会被放大，所以需要对矩阵B的方差进行缩放，以保证A*B的方差和A的方差相比不会被放大很多。

            if hasattr(module, 'NANOGPT_SCALE_INIT'):

                std *= (2 * self.config.n_layer) ** -0.5 #经过n次的x=x+y，x的方差会逐渐增加。因此需要根据残差流的层数或次数对权重的std进行缩放，乘以2是因此每个layer经过了2次残差连接，一次是atten，一次是mlp，对std进行缩放，以保证tensor x的方差不会随着残差流层数的增加再变大。

            torch.nn.init.normal_(module.weight, mean=0.0, std=std)

            if module.bias is not None:

                torch.nn.init.zeros_(module.bias) #bias初始化为0，pytorch默认是均匀分布

        elif isinstance(module, nn.Embedding):

            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

 

    def forward(self, idx, targets=None):

        # idx is of shape (B, T)

        B, T = idx.size()

        assert T <= self.config.block_size, f"Cannot forward sequence of length {T}, block size is only {self.config.block_size}"

        # forward the token and posisition embeddings

        pos = torch.arange(0, T, dtype=torch.long, device=idx.device) # shape (T)

        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (T, n_embd)，Batch中所有样本的position embeddings都是一样的

        tok_emb = self.transformer.wte(idx) # token embeddings of shape (B, T, n_embd)

        x = tok_emb + pos_emb #这里相加时其实进行了广播

        # forward the blocks of the transformer

        for block in self.transformer.h:

            x = block(x)

        # forward the final layernorm and the classifier

        x = self.transformer.ln_f(x)

        logits = self.lm_head(x) # (B, T, vocab_size)

        loss = None

        if targets is not None:

            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))#F.cross_entropy不支持多维输入，只支持2维输入，所以将logits(B, T, vocab_size)变为（B*T，vocab_size）二维tensor。targets要变为一维的。

        return logits, loss

 

    @classmethod

    def from_pretrained(cls, model_type):

        """Loads pretrained GPT-2 model weights from huggingface"""

        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}

        from transformers import GPT2LMHeadModel

        print("loading weights from pretrained gpt: %s" % model_type)

 

        # n_layer, n_head and n_embd are determined from model_type

        config_args = {

            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params

            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params

            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params

            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params

        }[model_type]

        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints

        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints

        # create a from-scratch initialized minGPT model

        config = GPTConfig(**config_args)

        model = GPT(config)

        sd = model.state_dict()

        sd_keys = sd.keys()

        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param

 

        # init a huggingface/transformers model

        model_hf = GPT2LMHeadModel.from_pretrained(model_type)

        sd_hf = model_hf.state_dict()

 

        # copy while ensuring all of the parameters are aligned and match in names and shapes

        sd_keys_hf = sd_hf.keys()

        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer

        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)

        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']

        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear

        # this means that we have to transpose these weights when we import them

        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"

        for k in sd_keys_hf:

            if any(k.endswith(w) for w in transposed):

                # special treatment for the Conv1D weights we need to transpose

                assert sd_hf[k].shape[::-1] == sd[k].shape

                with torch.no_grad():

                    sd[k].copy_(sd_hf[k].t())

            else:

                # vanilla copy over the other parameters

                assert sd_hf[k].shape == sd[k].shape

                with torch.no_grad():

                    sd[k].copy_(sd_hf[k])

 

        return model

 

    def configure_optimizers(self, weight_decay, learning_rate, device_type):

        # start with all of the candidate parameters (that require grad)

        param_dict = {pn: p for pn, p in self.named_parameters()}

        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}

        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.

        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.

        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]

        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]

        optim_groups = [

            {'params': decay_params, 'weight_decay': weight_decay},

            {'params': nodecay_params, 'weight_decay': 0.0}

        ]

        num_decay_params = sum(p.numel() for p in decay_params)

        num_nodecay_params = sum(p.numel() for p in nodecay_params)

        if master_process:

            print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")

            print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")

        # Create AdamW optimizer and use the fused version if it is available

        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters  #检查adamw的参数

        use_fused = fused_available and device_type == "cuda"

        if master_process:

            print(f"using fused AdamW: {use_fused}")

        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)

        return optimizer

 

# -----------------------------------------------------------------------------

import tiktoken

import numpy as np

 

def load_tokens(filename):

    npt = np.load(filename)

    npt = npt.astype(np.int32) # added after video

    ptt = torch.tensor(npt, dtype=torch.long)

    return ptt

 

class DataLoaderLite:

    def __init__(self, B, T, process_rank, num_processes, split):

        self.B = B

        self.T = T

        self.process_rank = process_rank

        self.num_processes = num_processes

        assert split in {'train', 'val'}

 

        # get the shard filenames

        data_root = "edu_fineweb10B"

        shards = os.listdir(data_root)

        shards = [s for s in shards if split in s]

        shards = sorted(shards)

        shards = [os.path.join(data_root, s) for s in shards]

        self.shards = shards

        assert len(shards) > 0, f"no shards found for split {split}"

        if master_process:

            print(f"found {len(shards)} shards for split {split}")

        self.reset()

 

    def reset(self):

        # state, init at shard zero

        self.current_shard = 0

        self.tokens = load_tokens(self.shards[self.current_shard])

        self.current_position = self.B * self.T * self.process_rank

 

    def next_batch(self):

        B, T = self.B, self.T

        buf = self.tokens[self.current_position : self.current_position+B*T+1]

        x = (buf[:-1]).view(B, T) # inputs

        y = (buf[1:]).view(B, T) # targets

        # advance the position in the tensor

        self.current_position += B * T * self.num_processes #ddp时，每次取数据的跨度要乘以num_processes，进程数量

        # if loading the next batch would be out of bounds, advance to next shard

        if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens):

            self.current_shard = (self.current_shard + 1) % len(self.shards)

            self.tokens = load_tokens(self.shards[self.current_shard])

            self.current_position = B * T * self.process_rank

        return x, y

 

# -----------------------------------------------------------------------------

# helper function for HellaSwag eval

# takes tokens, mask, and logits, returns the index of the completion with the lowest loss

 

def get_most_likely_row(tokens, mask, logits):

    # evaluate the autoregressive loss at all positions

    shift_logits = (logits[..., :-1, :]).contiguous()

    shift_tokens = (tokens[..., 1:]).contiguous()

    flat_shift_logits = shift_logits.view(-1, shift_logits.size(-1))

    flat_shift_tokens = shift_tokens.view(-1)

    shift_losses = F.cross_entropy(flat_shift_logits, flat_shift_tokens, reduction='none')

    shift_losses = shift_losses.view(tokens.size(0), -1)

    # now get the average loss just for the completion region (where mask == 1), in each row

    shift_mask = (mask[..., 1:]).contiguous() # we must shift mask, so we start at the last prompt token

    masked_shift_losses = shift_losses * shift_mask

    # sum and divide by the number of 1s in the mask

    sum_loss = masked_shift_losses.sum(dim=1)

    avg_loss = sum_loss / shift_mask.sum(dim=1)

    # now we have a loss for each of the 4 completions

    # the one with the lowest loss should be the most likely

    pred_norm = avg_loss.argmin().item()

    return pred_norm

 

# -----------------------------------------------------------------------------

# simple launch:

# python train_gpt2.py

# DDP launch for e.g. 8 GPUs:

# torchrun --standalone --nproc_per_node=8 train_gpt2.py

 

# run the training loop

from torch.distributed import init_process_group, destroy_process_group

from torch.nn.parallel import DistributedDataParallel as DDP

import torch.distributed as dist

 

# set up DDP (distributed data parallel).

# torchrun command sets the env variables RANK, LOCAL_RANK, and WORLD_SIZE

# 8块GPU，开8个进程，每个进程一个GPU，有一个主进程，rank=0，其他是辅助进程。

ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?

if ddp:

    # use of DDP atm demands CUDA, we set the device appropriately according to rank

    assert torch.cuda.is_available(), "for now i think we need CUDA for DDP"

    init_process_group(backend='nccl')

    ddp_rank = int(os.environ['RANK'])

    ddp_local_rank = int(os.environ['LOCAL_RANK']) # local_rank，只有在多节点时才使用，是单个 节点上的GPU的rank

    ddp_world_size = int(os.environ['WORLD_SIZE'])

    device = f'cuda:{ddp_local_rank}'

    torch.cuda.set_device(device)

    master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.

else:

    # vanilla, non-DDP run

    ddp_rank = 0

    ddp_local_rank = 0

    ddp_world_size = 1

    master_process = True

    # attempt to autodetect device

    device = "cpu"

    if torch.cuda.is_available():

        device = "cuda"

    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():

        device = "mps"

    print(f"using device: {device}")

 

# added after video, pytorch can be serious about it's device vs. device_type distinction

device_type = "cuda" if device.startswith("cuda") else "cpu"

 

torch.manual_seed(1337)

if torch.cuda.is_available():

    torch.cuda.manual_seed(1337)

 

enc = tiktoken.get_encoding("gpt2")

 

total_batch_size = 524288 # 2**19, ~0.5M, in number of tokens

B = 64 # micro batch size

T = 1024 # sequence length

assert total_batch_size % (B * T * ddp_world_size) == 0, "make sure total_batch_size is divisible by B * T * ddp_world_size"

grad_accum_steps = total_batch_size // (B * T * ddp_world_size) #grad_accum_steps要相应调整，因为ddp

if master_process: #只有主进程负责打印

    print(f"total desired batch size: {total_batch_size}")

    print(f"=> calculated gradient accumulation steps: {grad_accum_steps}")

 

train_loader = DataLoaderLite(B=B, T=T, process_rank=ddp_rank, num_processes=ddp_world_size, split="train")

val_loader = DataLoaderLite(B=B, T=T, process_rank=ddp_rank, num_processes=ddp_world_size, split="val")

 

torch.set_float32_matmul_precision('high')

 

# create model

model = GPT(GPTConfig(vocab_size=50304))

# model = GPT.from_pretrained("gpt2") # or init from OpenAI GPT-2

model.to(device) #model.to(device)将model转移到device上，但是对于tensor x来说，必须要用x=x.to(device)，因为x是没有状态的，x.to(device)不会将x转换为device上的对象，而是返回指向device上新内存的指针，。

use_compile = False # torch.compile interferes with HellaSwag eval and Generation. TODO fix

if use_compile:

    model = torch.compile(model)

if ddp:

    model = DDP(model, device_ids=[ddp_local_rank])

raw_model = model.module if ddp else model # always contains the "raw" unwrapped model 如果用ddp，那么raw_model是model.module

 

max_lr = 6e-4

min_lr = max_lr * 0.1

warmup_steps = 715

max_steps = 19073 # 19,073 steps is ~1 epoch, if data is 10B tokens and batch size 0.5M tokens

def get_lr(it):

    # 1) linear warmup for warmup_iters steps

    if it < warmup_steps:

        return max_lr * (it+1) / warmup_steps

    # 2) if it > lr_decay_iters, return min learning rate

    if it > max_steps:

        return min_lr

    # 3) in between, use cosine decay down to min learning rate

    decay_ratio = (it - warmup_steps) / (max_steps - warmup_steps)

    assert 0 <= decay_ratio <= 1

    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff starts at 1 and goes to 0

    return min_lr + coeff * (max_lr - min_lr)

 

# optimize!

optimizer = raw_model.configure_optimizers(weight_decay=0.1, learning_rate=6e-4, device_type=device_type)

 

# create the log directory we will write checkpoints to and log to

log_dir = "log"

os.makedirs(log_dir, exist_ok=True)

log_file = os.path.join(log_dir, f"log.txt")

with open(log_file, "w") as f: # open for writing to clear the file

    pass

 

for step in range(max_steps):

    t0 = time.time()

    last_step = (step == max_steps - 1)

 

    # once in a while evaluate our validation loss

    if step % 250 == 0 or last_step:

        model.eval()

        val_loader.reset()

        with torch.no_grad():

            val_loss_accum = 0.0

            val_loss_steps = 20

            for _ in range(val_loss_steps):

                x, y = val_loader.next_batch()

                x, y = x.to(device), y.to(device)

                with torch.autocast(device_type=device_type, dtype=torch.bfloat16):

                    logits, loss = model(x, y)

                loss = loss / val_loss_steps

                val_loss_accum += loss.detach()

        if ddp:

            dist.all_reduce(val_loss_accum, op=dist.ReduceOp.AVG)

        if master_process:

            print(f"validation loss: {val_loss_accum.item():.4f}")

            with open(log_file, "a") as f:

                f.write(f"{step} val {val_loss_accum.item():.4f}\n")

            if step > 0 and (step % 5000 == 0 or last_step):

                # optionally write model checkpoints

                checkpoint_path = os.path.join(log_dir, f"model_{step:05d}.pt")

                checkpoint = {

                    'model': raw_model.state_dict(),

                    'config': raw_model.config,

                    'step': step,

                    'val_loss': val_loss_accum.item()

                }

                # you might also want to add optimizer.state_dict() and

                # rng seeds etc., if you wanted to more exactly resume training

                torch.save(checkpoint, checkpoint_path)

 

    # once in a while evaluate hellaswag

    if (step % 250 == 0 or last_step) and (not use_compile):

        num_correct_norm = 0

        num_total = 0

        for i, example in enumerate(iterate_examples("val")):

            # only process examples where i % ddp_world_size == ddp_rank

            if i % ddp_world_size != ddp_rank:

                continue

            # render the example into tokens and labels

            _, tokens, mask, label = render_example(example)

            tokens = tokens.to(device)

            mask = mask.to(device)

            # get the logits

            with torch.no_grad():

                with torch.autocast(device_type=device_type, dtype=torch.bfloat16):

                    logits, loss = model(tokens)

                pred_norm = get_most_likely_row(tokens, mask, logits)

            num_total += 1

            num_correct_norm += int(pred_norm == label)

        # reduce the stats across all processes

        if ddp:

            num_total = torch.tensor(num_total, dtype=torch.long, device=device)

            num_correct_norm = torch.tensor(num_correct_norm, dtype=torch.long, device=device)

            dist.all_reduce(num_total, op=dist.ReduceOp.SUM) #rank之间进行同步

            dist.all_reduce(num_correct_norm, op=dist.ReduceOp.SUM) #rank之间进行同步

            num_total = num_total.item()

            num_correct_norm = num_correct_norm.item()

        acc_norm = num_correct_norm / num_total

        if master_process:

            print(f"HellaSwag accuracy: {num_correct_norm}/{num_total}={acc_norm:.4f}")

            with open(log_file, "a") as f:

                f.write(f"{step} hella {acc_norm:.4f}\n")

 

    # once in a while generate from the model (except step 0, which is noise)

    if ((step > 0 and step % 250 == 0) or last_step) and (not use_compile):

        model.eval()

        num_return_sequences = 4

        max_length = 32

        tokens = enc.encode("Hello, I'm a language model,")

        tokens = torch.tensor(tokens, dtype=torch.long)

        tokens = tokens.unsqueeze(0).repeat(num_return_sequences, 1)

        xgen = tokens.to(device)

        sample_rng = torch.Generator(device=device)

        sample_rng.manual_seed(42 + ddp_rank)

        while xgen.size(1) < max_length:

            # forward the model to get the logits

            with torch.no_grad():

                with torch.autocast(device_type=device_type, dtype=torch.bfloat16):

                    logits, loss = model(xgen) # (B, T, vocab_size)

                # take the logits at the last position

                logits = logits[:, -1, :] # (B, vocab_size)

                # get the probabilities

                probs = F.softmax(logits, dim=-1)

                # do top-k sampling of 50 (huggingface pipeline default)

                # topk_probs here becomes (5, 50), topk_indices is (5, 50)

                topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)

                # select a token from the top-k probabilities

                # note: multinomial does not demand the input to sum to 1

                ix = torch.multinomial(topk_probs, 1, generator=sample_rng) # (B, 1)

                # gather the corresponding indices

                xcol = torch.gather(topk_indices, -1, ix) # (B, 1)

                # append to the sequence

                xgen = torch.cat((xgen, xcol), dim=1)

        # print the generated text

        for i in range(num_return_sequences):

            tokens = xgen[i, :max_length].tolist()

            decoded = enc.decode(tokens)

            print(f"rank {ddp_rank} sample {i}: {decoded}")

 

    # do one step of the optimization 这里开始训练

    model.train()

    optimizer.zero_grad()

    loss_accum = 0.0

    for micro_step in range(grad_accum_steps): #梯度累积步数，在每个step中，只进行loss.backward，只进行梯度累加，不进行梯度更新

        x, y = train_loader.next_batch()

        x, y = x.to(device), y.to(device)

        # added after video, this field is also used by the forward pass.

        if ddp:

            model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1) #loss.backward()会进行梯度同步 ，在这里设置只有在grad_accum_steps的最后一步时，再进行所有rank的梯度同步。

        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):#进行bf16运算。但并不是model中的所有层都会被转换为bf16，只有部分层会进行bf16转换，其他层如softmax等仍然保持在fp32，因为这些层对精度变化更敏感。矩阵乘法对精度的变化 相对来说很稳健。

            logits, loss = model(x, y)

        # we have to scale the loss to account for gradient accumulation,

        # because the gradients just add on each successive backward().

        # addition of gradients corresponds to a SUM in the objective, but

        # instead of a SUM we want MEAN. Scale the loss here so it comes out right

        loss = loss / grad_accum_steps

        loss_accum += loss.detach()

        loss.backward()

    if ddp:

        dist.all_reduce(loss_accum, op=dist.ReduceOp.AVG) #这行代码对所有rank上的loss进行平均，并同步存储在每个rank上

    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    # determine and set the learning rate for this iteration

    lr = get_lr(step)

    for param_group in optimizer.param_groups:

        param_group['lr'] = lr

    optimizer.step()

    if device_type == "cuda":

        torch.cuda.synchronize() # wait for the GPU to finish work

    t1 = time.time()

    dt = t1 - t0 # time difference in seconds

    tokens_processed = train_loader.B * train_loader.T * grad_accum_steps * ddp_world_size

    tokens_per_sec = tokens_processed / dt

    if master_process:

        print(f"step {step:5d} | loss: {loss_accum.item():.6f} | lr {lr:.4e} | norm: {norm:.4f} | dt: {dt*1000:.2f}ms | tok/sec: {tokens_per_sec:.2f}")

        with open(log_file, "a") as f:

            f.write(f"{step} train {loss_accum.item():.6f}\n")

 

if ddp:

    destroy_process_group()